import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import welch, detrend

def fetch_futures_data(ticker, start_date, end_date):
    """从本地CSV文件获取期货数据"""
    # 根据ticker选择对应的本地文件
    file_map = {
        'GC=F': 'data/SA0SH.csv',  # 黄金期货
        'SI=F': 'data/SF0SH.csv'   # 白银期货
    }
    
    # 读取数据
    data = pd.read_csv(file_map[ticker],index_col='date_time')
    
    # 仅保留收盘价
    data = data[['close']]
    
    # 按日期范围筛选
    data = data.loc[start_date:end_date]
    
    return data

def preprocess_data(data):
    """数据预处理：去除NaN值，去趋势"""
    # 填补缺失值
    data = data.fillna(method='ffill').fillna(method='bfill')
    
    # 去趋势
    data_detrended = data.apply(detrend, axis=0)
    
    return data_detrended

def calculate_psd(data, fs=1.0, nperseg=256):
    """计算功率谱密度"""
    # 确保输入是一维数组
    if len(data.shape) > 1:
        data = data.values.flatten()
    
    # 计算PSD
    freqs, psd = welch(data, fs=fs, nperseg=min(nperseg, len(data)), scaling='density')
    
    return freqs, psd

def plot_psd(freqs, psd, title):
    """绘制PSD图"""
    plt.figure(figsize=(12, 6))
    
    # 绘制功率谱密度
    plt.subplot(1, 2, 1)
    plt.semilogy(freqs, psd)
    plt.xlabel('Frequency')
    plt.ylabel('Power Spectral Density')
    plt.title(title)
    plt.grid(alpha=0.3)
    
    # 绘制对数-对数图
    plt.subplot(1, 2, 2)
    plt.loglog(freqs, psd)
    plt.xlabel('Frequency')
    plt.ylabel('Power Spectral Density')
    plt.title(title + ' (Log-Log Scale)')
    plt.grid(alpha=0.3)
    
    plt.tight_layout()
    plt.show()

def main():
    # 参数设置
    ticker = 'GC=F'  # 黄金期货代码
    start_date = '2010-01-01'
    end_date = '2023-12-31'
    
    # 获取数据
    print("获取数据...")
    data = fetch_futures_data(ticker, start_date, end_date)
    print(f"原始数据形状: {data.shape}")
    
    # 数据预处理
    print("数据预处理...")
    data_detrended = preprocess_data(data)
    print(f"预处理后数据形状: {data_detrended.shape}")
    
    # 计算PSD
    print("计算功率谱密度...")
    freqs, psd = calculate_psd(data_detrended)
    print(f"频率数据形状: {freqs.shape}, PSD数据形状: {psd.shape}")
    
    # 可视化结果
    print("绘制结果...")
    plot_psd(freqs, psd, f'Futures Price PSD ({ticker})')
    
    # 打印主要频率分量
    peak_indices = np.argsort(psd[::-1][:10])[::-1]  # 获取最大的10个频率
    peak_frequencies = freqs[peak_indices]
    
    print("\n主要频率分量:")
    for i, freq in enumerate(peak_frequencies):
        print(f"第{i+1}大频率: {freq:.6f} Hz")

if __name__ == "__main__":
    main()
